import pymysql
import bcrypt
from pymysql.cursors import DictCursor


class UserModel:
    def __init__(self, db_config):
        self.db_config = db_config

    def get_connection(self):
        return pymysql.connect(
            host=self.db_config['host'],
            user=self.db_config['user'],
            password=self.db_config['password'],
            database=self.db_config['database'],
            cursorclass=DictCursor
        )

    def authenticate(self, username, password):
        try:
            with self.get_connection() as conn:
                with conn.cursor() as cursor:
                    sql = "SELECT * FROM users WHERE username = %s"
                    cursor.execute(sql, (username,))
                    user = cursor.fetchone()

                    if user:
                        # 调试输出
                        print(f"Stored hash: {user['password']}")
                        print(f"Input password: {password}")

                        # 修复密码验证
                        if isinstance(user['password'], str):
                            if bcrypt.checkpw(password.encode('utf-8'), user['password'].encode('utf-8')):
                                return user
                        return None
        except Exception as e:
            print(f"Authentication error: {e}")
            return None

    def create_user(self, username, password):
        try:
            # 检查用户名是否已存在
            if self.get_user_by_username(username):
                raise Exception('用户名已存在')

            # 加密密码
            hashed_password = bcrypt.hashpw(password.encode('utf-8'), bcrypt.gensalt())

            with self.get_connection() as conn:
                with conn.cursor() as cursor:
                    sql = """
                    INSERT INTO users (username, password) 
                    VALUES (%s, %s)
                    """
                    # 确保存储的是字符串
                    cursor.execute(sql, (username, hashed_password.decode('utf-8')))
                    conn.commit()
        except pymysql.Error as e:
            raise Exception(f'数据库错误: {e}')

    def get_user_by_username(self, username):
        try:
            with self.get_connection() as conn:
                with conn.cursor() as cursor:
                    sql = "SELECT * FROM users WHERE username = %s"
                    cursor.execute(sql, (username,))
                    return cursor.fetchone()
        except pymysql.Error as e:
            print(f"Database error: {e}")
            return None